from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer  # 导入 AutoTokenizer
import torch

# 🔹 基础模型 & LoRA 适配器路径
base_model_name = "Qwen/Qwen1.5-4B-Chat"
lora_path = "./results/checkpoint-825"

# 🔹 加载未量化的基础模型到 CPU
model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    device_map=None,  # 不使用 device_map
    torch_dtype=torch.float16,  # 使用 FP32 精度
).to(torch.device("cpu"))  # 显式加载到 CPU

# 🔹 加载分词器
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

# 🔹 加载 LoRA 适配器
model = PeftModel.from_pretrained(model, lora_path)

# 🔹 合并 LoRA 权重
model = model.merge_and_unload()

# 🔹 保存合并后的完整模型和分词器
save_path = "./qwen1.5-4b-chat-xiaoyu-merged"
model.save_pretrained(save_path)
tokenizer.save_pretrained(save_path)  # 保存分词器

print(f"✅ LoRA 已合并，模型和分词器保存在: {save_path}")
